import numpy as np
from environment.trace_sample import Trace
from environment.state import State
from environment.env import Environment
from learning.eval_helper import Measurement
from config import DRLParameters

measurement = Measurement()


def measure_result(dqn_model, parameter_spaces: DRLParameters, training: bool, alpha, points=999):
    """ 评估DQN模型
    :param alpha:
    :param dqn_model: 模型
    :param parameter_spaces: 参数空间
    :param training: 训练模式
    :param points:
    :return:
    """
    if parameter_spaces.DEBUG:
        case_ids = np.array(parameter_spaces.INITIAL_CASE_IDS)
    elif training:
        case_ids = Trace(parameter_spaces).train_case_ids
    else:
        case_ids = Trace(parameter_spaces).test_case_ids

    cnt = 0
    for cid in case_ids:
        cnt += 1
        trace = Trace(parameter_space=parameter_spaces, case_id=cid)
        env = Environment(trace=trace, parameter_space=parameter_spaces, training=training, num_points=points)
        s: State = env.state
        while True:
            action = dqn_model.choose_action(s.feature)
            next_state, _, done, _ = env.step(action)
            if done:
                x, ypred, ytrue = measurement.prepare_x_y(next_state)
                loss = measurement.AILoss(x, ypred, ytrue, alpha=alpha)
                measurement.collect_result(loss, len(next_state.actions))
                break
            s = next_state
    measurement.display_result()


if __name__ == '__main__':
    pass
